/**
 * @date 2016年3月6日 下午8:42:36
 * @version V1.0
 */
package edu.cn.nlsde.tmfst.lda;

import edu.cn.nlsde.Utils.IOUtil;
import edu.cn.nlsde.tmfst.interfaces.AbstractLDA;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

//import edu.cn.nlsde.tmfst.processor.data_pre_process_Chinese;

/**
 * @author ffftzh
 * @ClassName: LDA_WN
 * @Description: LDA_WN的实现
 * @date 2016年3月6日 下午8:42:36
 */
public class LDA_WN extends AbstractLDA implements Runnable {

    private int[] word_sum;
    private HashMap<Integer, HashMap<Integer, Integer>> word_word_count = new HashMap<>();

    public LDA_WN(List<List<String>> data, int save_step, int topic_num, int iter_num, double alpha, double beta) {

        this.alpha = alpha;
        this.beta = beta;
        this.save_step = save_step;
        this.iter_num = iter_num;
        this.topic_num = topic_num;
        this.data_path = "";
        this.outputPath = "tmp/LDA_WN_" + alpha + "_" + beta + "_" + topic_num + "/";
        (new File(this.outputPath)).mkdirs();
        print_parameter();
        this.load_data(data);
    }

    public LDA_WN(String data_path, int topic_num, int iter_num, double alpha, double beta, int threadId) {

        this.alpha = alpha;
        this.beta = beta;
        this.iter_num = iter_num;
        this.topic_num = topic_num;
        this.data_path = data_path;
        this.outputPath = (new File(data_path)).getParentFile().getAbsolutePath() + "/LDA_WN_" + alpha + "_" + beta + "_" + topic_num + "/" + threadId + "/";
        (new File(this.outputPath)).mkdirs();
        print_parameter();
        this.load_data_from_file();
    }

    public static void build_LDA_WN_Model(String data_path, int topic_num, int iter_num, double alpha, double beta, int threadNum) throws InterruptedException {

        ExecutorService service = Executors.newCachedThreadPool();
        for (int threadId = 0; threadId < threadNum; threadId++) {
            LDA_WN model = new LDA_WN(data_path, topic_num, iter_num, alpha, beta, threadId);
            service.submit(model);
        }
        service.shutdown();
        service.awaitTermination(10, TimeUnit.DAYS);


//        LDA_WN model = new LDA_WN(data_path, topic_num, iter_num, alpha, beta, 0);
//        model.run();

    }

    private void print_parameter() {
        System.out.println("method:LDA_WN"
                + "	path:" + this.data_path
                + "\talpha:" + this.alpha
                + "\tbeta:" + this.beta
                + "\titer_num:" + this.iter_num
                + "\ttopic_num:" + this.topic_num
        );
    }

    private void load_data(List<List<String>> data) {

        this.wordId_of_corpus = new int[data.size()][];
        int docIndex = 0;
        for (List<String> words : data) {
            this.wordId_of_corpus[docIndex] = new int[words.size()];
            int wordIndex = 0;
            for (String oneWord : words) {
                if (!this.wordMap.containsKey(oneWord)) {
                    this.wordMap.put(oneWord, this.wordMap.size());
                    this.id2Word.put(this.wordMap.get(oneWord), oneWord);
                }
                this.wordId_of_corpus[docIndex][wordIndex] = this.wordMap.get(oneWord);
                wordIndex++;
            }
            docIndex++;
        }
    }

    private void load_data_from_file() {
        try {
            BufferedReader reader = IOUtil.getReader(this.data_path);
            String line;
            ArrayList<int[]> tmpCorpus = new ArrayList<>();
            while ((line = reader.readLine()) != null) {
                String[] words = line.split("\\s+");
                int[] oneDoc = new int[words.length];
                for (int wordIndex = 0; wordIndex < words.length; wordIndex++) {
                    if (!this.wordMap.containsKey(words[wordIndex])) {
                        this.wordMap.put(words[wordIndex], this.wordMap.size());
                        this.id2Word.put(this.wordMap.get(words[wordIndex]), words[wordIndex]);
                    }
                    oneDoc[wordIndex] = this.wordMap.get(words[wordIndex]);
                }
                tmpCorpus.add(oneDoc);
            }
            reader.close();

            this.wordId_of_corpus = new int[tmpCorpus.size()][];
            for (int docIndex = 0; docIndex < this.wordId_of_corpus.length; docIndex++) {
                this.wordId_of_corpus[docIndex] = tmpCorpus.get(docIndex);
            }

        } catch (Exception e) {
            e.printStackTrace();
        }

    }
//
//    private void save_twords(int topWordNum) throws IOException {
//        BufferedWriter writer = IOUtil.getWriter(this.outputPath + "model-final.twords");
//
//        for (int topic_index = 0; topic_index < this.word_num_in_topic_word.length; topic_index++) {
//            HashMap<Integer, Double> oneLine = new HashMap<>();
//            for (int word_id = 0; word_id < this.word_num_in_topic_word[topic_index].length; word_id++) {
//                oneLine.put(word_id, ((double) this.word_num_in_topic_word[topic_index][word_id]) / this.word_num_in_topic[topic_index]);
//            }
//            List<Map.Entry<Integer, Double>> maplist = new ArrayList<>(oneLine.entrySet());
//
//            Collections.sort(maplist, (o1, o2) -> o2.getValue().compareTo(o1.getValue()));
//
//            writer.write("Topic:" + topic_index + "\n");
//            int count = 0;
//            for (Map.Entry<Integer, Double> o1 : maplist) {
//                writer.write("\t" + this.id2Word.get(o1.getKey()) + "\t:" + o1.getValue() + "\n");
//                count++;
//                if (count > topWordNum) {
//                    break;
//                }
//            }
//            writer.write("\n");
//
//        }
//
//
//        writer.close();
//    }

    private void init_model() {
        this.topicId_of_corpus = new int[this.wordId_of_corpus.length][];

        this.word_num_in_topic_word = new int[this.topic_num][this.wordMap.size()];
        this.word_num_in_topic = new int[this.topic_num];
        this.word_num_in_doc_topic = new int[this.wordId_of_corpus.length][this.topic_num];
        this.word_num_in_doc = new int[this.wordId_of_corpus.length];

        this.word_sum = new int[this.wordMap.size()];

        for (int docIndex = 0; docIndex < this.wordId_of_corpus.length; docIndex++) {
            this.topicId_of_corpus[docIndex] = new int[this.wordId_of_corpus[docIndex].length];
            this.word_num_in_doc[docIndex] = this.wordId_of_corpus[docIndex].length;
            for (int wordIndex = 0; wordIndex < this.wordId_of_corpus[docIndex].length; wordIndex++) {
                int wordId = this.wordId_of_corpus[docIndex][wordIndex];
                int topicId = (int) (Math.random() * this.topic_num);

                this.topicId_of_corpus[docIndex][wordIndex] = topicId;

                this.word_num_in_topic_word[topicId][wordId] += 1;
                this.word_num_in_topic[topicId] += 1;
                this.word_num_in_doc_topic[docIndex][topicId] += 1;

                for (int wordIndex2 = 0; wordIndex2 < this.wordId_of_corpus[docIndex].length; wordIndex2++) {
                    if (this.wordId_of_corpus[docIndex][wordIndex] == this.wordId_of_corpus[docIndex][wordIndex2]) {
                        continue;
                    }
                    HashMap<Integer, Integer> theWord = word_word_count.get(this.wordId_of_corpus[docIndex][wordIndex]);
                    if (theWord == null) {
                        theWord = new HashMap<>();
                        word_word_count.put(this.wordId_of_corpus[docIndex][wordIndex], theWord);
                    }
                    Integer theNum = theWord.get(this.wordId_of_corpus[docIndex][wordIndex2]);
                    if (theNum == null) {
                        theNum = 0;
                    }
                    theWord.put(this.wordId_of_corpus[docIndex][wordIndex2], theNum + 1);
                    word_sum[this.wordId_of_corpus[docIndex][wordIndex]] += 1;
                }

            }
        }
    }

    private void save_wordMap() throws IOException {
        BufferedWriter writer = IOUtil.getWriter(this.outputPath + "model-final.wordmap");
        for (String key : this.wordMap.keySet()) {
            writer.write(key + " " + this.wordMap.get(key));
            writer.write("\n");
        }
        writer.close();
    }

    private void save_theta(int saveStep) throws IOException {
        BufferedWriter writer = IOUtil.getWriter(this.outputPath + "model-" + saveStep + ".theta");
//        BufferedWriter writer2 = IOUtil.getWriter(this.outputPath + "model-final.theta2");
        int doc_index = 0;
        for (int[] line : this.word_num_in_doc_topic) {
            for (int topic : line) {
                writer.write((this.alpha + topic) / (this.word_num_in_doc[doc_index] + this.alpha * this.topic_num) + " ");
//                writer2.write(((double)topic) / (this.word_num_in_doc[doc_index]) + " ");
            }
            writer.write("\n");
//            writer2.write("\n");
            doc_index++;
        }
        writer.close();
//        writer2.close();
    }

    private void save_tassign(int saveStep) throws IOException {
        BufferedWriter writer = IOUtil.getWriter(this.outputPath + "model-" + saveStep + ".tassign");
        for (int doc_index = 0; doc_index < this.topicId_of_corpus.length; doc_index++) {
            for (int word_index = 0; word_index < this.topicId_of_corpus[doc_index].length; word_index++) {
                writer.write(this.wordId_of_corpus[doc_index][word_index] + ":" + this.topicId_of_corpus[doc_index][word_index] + " ");
            }
            writer.write("\n");
        }
        writer.close();
    }

    private void save_phi(int saveStep) throws IOException {
        BufferedWriter writer = IOUtil.getWriter(this.outputPath + "model-" + saveStep + ".phi");
//        BufferedWriter writer2 = IOUtil.getWriter(this.outputPath + "model-final.phi2");
        int topic_index = 0;
        for (int[] line : this.word_num_in_topic_word) {
            for (int words : line) {
                writer.write((this.beta + words) / (this.word_num_in_topic[topic_index] + this.beta * this.wordMap.size()) + " ");
//                writer2.write((1.0*words) / (this.word_num_in_topic[topic_index]) + " ");
            }
            writer.write("\n");
//            writer2.write("\n");
            topic_index++;
        }
        writer.close();
//        writer2.close();

    }

    private int sample(int docIndex, int wordIndex) {
        int wordId = this.wordId_of_corpus[docIndex][wordIndex];
        int oldTopicId = this.topicId_of_corpus[docIndex][wordIndex];
        int newTopicId = -1;

        this.word_num_in_topic_word[oldTopicId][wordId] -= 1;
        this.word_num_in_topic[oldTopicId] -= 1;
        this.word_num_in_doc_topic[docIndex][oldTopicId] -= 1;

        double[] p = new double[this.topic_num];
        for (int k = 0; k < this.topic_num; k++) {
            p[k] = 0;
            for (Integer key : word_word_count.get(wordId).keySet()) {
                p[k] = (word_num_in_doc_topic[docIndex][k] + alpha)
                        * (word_num_in_topic_word[k][key] + beta)
                        / (word_num_in_topic[k] + word_num_in_topic_word[0].length * beta)
                        * word_word_count.get(wordId).get(key)
                        / word_sum[key];
            }

        }

        for (int k = 1; k < this.topic_num; k++) {
            p[k] += p[k - 1];
        }

        double u = Math.random() * p[this.topic_num - 1];
        for (int k = 0; k < this.topic_num; k++) {
            if (u < p[k]) {
                newTopicId = k;
                break;
            }

        }

        this.word_num_in_topic_word[newTopicId][wordId] += 1;
        this.word_num_in_topic[newTopicId] += 1;
        this.word_num_in_doc_topic[docIndex][newTopicId] += 1;

        assert newTopicId > -1;
        return newTopicId;

    }

    private void build_model() {

        for (int iter = 0; iter < this.iter_num; iter++) {
            long startTime = System.currentTimeMillis();

            for (int docIndex = 0; docIndex < this.wordId_of_corpus.length; docIndex++) {
                for (int wordIndex = 0; wordIndex < this.wordId_of_corpus[docIndex].length; wordIndex++) {
                    this.topicId_of_corpus[docIndex][wordIndex] = this.sample(docIndex, wordIndex);
                }
            }

            if (iter % save_step == 0) {
                save_result(iter);
            }

            System.out.println("finished iter :" + iter + "\tcost time:" + ((double) System.currentTimeMillis() - startTime) / 1000);
        }

    }

    private void save_result(int saveStep) {
        try {
            this.save_twords(20, saveStep);
            this.save_theta(saveStep);
            this.save_tassign(saveStep);
            this.save_wordMap();
            this.save_phi(saveStep);

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @Override
    public void run() {
        this.init_model();
        this.build_model();
        this.save_result(2000);
    }


}
